import torch
import torch.nn as nn
import torchvision.models as models


class ResNetEncoder(nn.Module):
    """ResNet作为编码器"""
    
    def __init__(self, pretrained=True, model_name='resnet50'):
        super(ResNetEncoder, self).__init__()
        
        if model_name == 'resnet50':
            resnet = models.resnet50(pretrained=pretrained)
        elif model_name == 'resnet34':
            resnet = models.resnet34(pretrained=pretrained)
        elif model_name == 'resnet18':
            resnet = models.resnet18(pretrained=pretrained)
        else:
            raise ValueError(f"Unsupported ResNet model: {model_name}")
        
        # 提取ResNet的各个层
        self.conv1 = resnet.conv1
        self.bn1 = resnet.bn1
        self.relu = resnet.relu
        self.maxpool = resnet.maxpool
        
        self.layer1 = resnet.layer1
        self.layer2 = resnet.layer2
        self.layer3 = resnet.layer3
        self.layer4 = resnet.layer4
        
        # 冻结预训练参数
        if pretrained:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        enc1 = x
        
        x = self.maxpool(x)
        enc2 = self.layer1(x)
        enc3 = self.layer2(enc2)
        enc4 = self.layer3(enc3)
        enc5 = self.layer4(enc4)
        
        return [enc1, enc2, enc3, enc4, enc5]


class DecoderBlock(nn.Module):
    """解码器块"""
    
    def __init__(self, in_channels, out_channels):
        super(DecoderBlock, self).__init__()
        self.conv = nn.Sequential(
            nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.conv(x)


class ResNetSegmentation(nn.Module):
    """基于ResNet的分割模型"""
    
    def __init__(self, n_classes=2, pretrained=True, model_name='resnet50'):
        super(ResNetSegmentation, self).__init__()
        self.encoder = ResNetEncoder(pretrained=pretrained, model_name=model_name)
        
        # 根据不同的ResNet模型设置通道数
        if model_name in ['resnet50', 'resnet101', 'resnet152']:
            channels = [2048, 1024, 512, 256, 64]
        else:  # resnet18, resnet34
            channels = [512, 256, 128, 64, 64]
        
        # 解码器
        self.decoder5 = DecoderBlock(channels[0], channels[1])
        self.decoder4 = DecoderBlock(channels[1], channels[2])
        self.decoder3 = DecoderBlock(channels[2], channels[3])
        self.decoder2 = DecoderBlock(channels[3], channels[4])
        
        # 最终分类层
        self.final_conv = nn.Sequential(
            nn.ConvTranspose2d(channels[4], 32, kernel_size=2, stride=2),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, n_classes, kernel_size=1)
        )

    def forward(self, x):
        # 编码
        features = self.encoder(x)
        
        # 解码
        d5 = self.decoder5(features[4])
        d4 = self.decoder4(d5)
        d3 = self.decoder3(d4)
        d2 = self.decoder2(d3)
        
        # 最终输出
        out = self.final_conv(d2)
        return out


def get_resnet_segmentation(n_classes=2, pretrained=True, model_name='resnet50'):
    """获取ResNet分割模型实例"""
    return ResNetSegmentation(n_classes=n_classes, pretrained=pretrained, model_name=model_name)


if __name__ == "__main__":
    # 测试模型
    model = get_resnet_segmentation(n_classes=2, pretrained=False, model_name='resnet50')
    x = torch.randn(1, 3, 256, 256)
    output = model(x)
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")
    print(f"模型参数量: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")

